# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 2022/1/16 17:06
# @Author : Wang Zhixing
# @File : Out2File.py
# @Software: PyCharm

from sklearn.cluster import KMeans

from ProcessData.PrintFormat.ColorPrint import GreenPrint
from .output_mehod.Vote import Vote
from .output_mehod.average_symbol2file import average_symbol2file
from .output_mehod.result2rsf_file import result2rsf_file


def Out2File(out=None,dataset=None,preds=None,**kwarg):
    if kwarg["model_name"]=="random":
        result2rsf_file(kwarg["root"], preds, kwarg["outfile_path"])
    elif kwarg["data_type"]=="symbol":
        if kwarg["out_method"] == "vote":
            vote = Vote(kwarg["root"])
            new_pred = vote.vote2pred(dataset.dic, preds)
            print("num of file after vote:" + str(len(new_pred)))
            vote.predtofile(kwarg["outfile_path"])
        elif kwarg["out_method"] == "average":
            filevector = average_symbol2file(out, dataset.dic, kwarg["root"], dim=kwarg["model_layer"][0])
            GreenPrint("将GNN训练之后的节点的vector平均到文件上，文件的数量为："+str(len(filevector)))
            kmeans_input = filevector
            kmeans = KMeans(n_clusters=kwarg["cluster"], random_state=0).fit(kmeans_input)
            preds = kmeans.predict(kmeans_input)
            result2rsf_file(kwarg["root"], preds, kwarg["outfile_path"])
        elif kwarg['out_method']=="fileaverage":
            result2rsf_file(kwarg["root"], preds, kwarg["outfile_path"])